use super::*;
use riir::agent;
use riir::agent::tools::playground::{run_rustycli, VerifyResp, VerifyType};
use std::marker::PhantomData;
use tracing::info;

use serde::{Deserialize, Serialize};

pub const DOCS_RUST_PATH: &str = "./files/rust/docs.md";

#[derive(Serialize, Deserialize)]
pub struct RustDocsInfo {
    pub rust_expert: String,
}

pub const RUST_DOCS: Prompt<RustDocsInfo> = Prompt("RustDocsInfo: [Rust Docs]", PhantomData);

#[derive(Serialize, Deserialize)]
pub struct DocsPlanInfo {
    pub task: String,
}

pub const RUST_DOCS_PLAN: Prompt<DocsPlanInfo> = Prompt(
    r#"
You are a Rust expert assistant. Your task is to add comments and documentation comments to the code, and nothing else.

Follow these rules when adding comments:

- Please ensure the code is fully functional and doesn't omit any details.
- Add documentation comments to items with pub visibility syntax.
- Add a // SAFETY comment to unsafe blocks to explain why they are safe.
- For unsafe functions, add /// # Safety in the documentation comments to explain the conditions for safely using this unsafe function.
- The Rust code snippet after your refactoring must not have any escape characters, please!


Please ensure that your reply strictly adheres to the following  JSON format:
```json
{{
	"plan": "<concise plan on how you will complete the task>",
	"code": "
	```rust
	<The Rust code after your audit, and the Rust code snippet without any escape characters Please!>
	```
	"
}}
```
"#,
    PhantomData,
);

#[derive(Debug)]
pub struct DocsAgent {
    attributes: AgentBase,
    bug_errors: Option<String>,
    bug_count: u8,
    retry_count: u8,
    llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
}

impl DocsAgent {
    pub fn new(llm: Arc<Mutex<Box<dyn LLM + 'static>>>) -> Self {
        let attributes: AgentBase = AgentBase {
            objective: "Docs Rust Code".to_string(),
            position: "Docs Agent".to_string(),
            state: AgentState::Starting,
            memory: vec![],
        };

        Self {
            attributes,
            bug_errors: None,
            bug_count: 0,
            retry_count: 0,
            llm: llm.clone(),
        }
    }

    async fn call_initial_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        // Concatenate Instruction
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_DOCS_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "docs", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_improved_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_DOCS_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "docs", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_fix_code_bugs(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_FIXED_EXPERT.0),
            format!(
                "CODE: {:?} \n ERROR_BUGS: {:?} \n",
                factsheet.code.as_ref().unwrap_or(&"".to_string()),
                self.bug_errors
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "docs", &ai_response).await;
        factsheet.code = Some(ai_response);
    }
}

#[async_trait]
impl AgentFns for DocsAgent {
    fn get_attributes_from_agent(&self) -> &AgentBase {
        &self.attributes
    }

    async fn execute(
        &mut self,
        factsheet: &mut FactSheet,
        output_path: PathBuf
    ) -> Result<(), Box<dyn std::error::Error>> {
        while self.attributes.state != AgentState::Done {
            match &self.attributes.state {
                AgentState::Starting => {
                    info!("Docs Agent Start ---");
                    self.call_initial_code(factsheet, output_path.clone()).await;
                    self.attributes.state = AgentState::Working;
                    continue;
                }

                AgentState::Working => {
                    info!("Docs Agent Working ---");
                    if self.bug_count == 0 {
                        self.call_improved_code(factsheet, output_path.clone()).await;
                    } else {
                        self.call_fix_code_bugs(factsheet, output_path.clone()).await;
                    }
                    self.attributes.state = AgentState::Debug;
                    continue;
                }

                AgentState::Debug => {
                    info!("Docs Agent Debug ---");

                    // --- 验证 LLM 响应代码
                    if let Some(e) = agent::capture_rust_code(&factsheet.code.clone().unwrap()) {
                        info!("EndANswer Chat: plain {:?}", e);

                        // do something async
                        let verft = run_rustycli(&e, VerifyType::BUILD).await;
                        info!("Verify Chat: {:?}", verft);
                        match verft {
                            VerifyResp::Success { stdout, stderr } => {
                                self.bug_count = 0;
                                info!("Code build successful...");
                                self.attributes.state = AgentState::Done;
                            }
                            VerifyResp::Fail(error_str) => {
                                // Update error stats
                                self.bug_count += 1;
                                self.retry_count += 1;
                                self.bug_errors = Some(error_str);
                                if self.retry_count > 10 {
                                    info!(
                                        "Code build failed but retry count is greater then 10... "
                                    );
                                    self.attributes.state = AgentState::Done;
                                } else {
                                    // Pass back for rework
                                    self.attributes.state = AgentState::Working;
                                    continue;
                                }
                            }
                        }
                    } else {
                        info!("Uncapture rust code!");
                        // Pass back for rework
                        self.attributes.state = AgentState::Working;
                        continue;
                    }
                }

                _ => {}
            }
        }
        Ok(())
    }
}
